#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Copyright (c) Huawei Technologies Co., Ltd. 2021. All rights reserved.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
import sys
import argparse
import logging
from ftrace_format_parser import FtraceEventCodeGenerator


class Common:
    ftrace_bundle_proto = "ftrace_event.proto"
    auto_generated_gni = "autogenerated.gni"

    this_file = os.path.basename(__file__)
    logging.basicConfig(
        format="%(asctime)s %(levelname)s %(message)s", level=logging.INFO
    )
    logger = logging.getLogger(this_file)

    proto_common_header = (
        ""
        "// THIS FILE IS GENERATED BY {}, PLEASE DON'T EDIT IT!\n"
        "// Copyright (c) Huawei Technologies Co., Ltd. 2021. All rights reserved.\n"
        '// Licensed under the Apache License, Version 2.0 (the "License");\n'
        "// you may not use this file except in compliance with the License.\n"
        "// You may obtain a copy of the License at\n"
        "//\n"
        "//     http://www.apache.org/licenses/LICENSE-2.0\n"
        "//\n"
        "// Unless required by applicable law or agreed to in writing, software\n"
        '// distributed under the License is distributed on an "AS IS" BASIS,\n'
        "// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n"
        "// See the License for the specific language governing permissions and\n"
        "// limitations under the License.\n"
        "//\n\n"
        'syntax = "proto3";\n\n'
        "option optimize_for = LITE_RUNTIME;\n\n".format(this_file)
    )

    gn_copyright_header = (
        ""
        "# THIS FILE IS GENERATE BY {}, PLEASE DON'T EDIT IT!\n"
        "# Copyright (c) Huawei Technologies Co., Ltd. 2021. All rights reserved.\n"
        '# Licensed under the Apache License, Version 2.0 (the "License");\n'
        "# you may not use this file except in compliance with the License.\n"
        "# You may obtain a copy of the License at\n"
        "#\n"
        "#     http://www.apache.org/licenses/LICENSE-2.0\n"
        "#\n"
        "# Unless required by applicable law or agreed to in writing, software\n"
        '# distributed under the License is distributed on an "AS IS" BASIS,\n'
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n"
        "# See the License for the specific language governing permissions and\n"
        "# limitations under the License.\n".format(this_file)
    )

    proto_message_head = (
        ""
        "message FtraceEvent {\n"
        "    uint64 timestamp = 1;\n"
        "    int32 tgid = 2;\n"
        "    string comm = 3;\n\n"
        "    message CommonFileds {\n"
        "        uint32 type = 1;\n"
        "        uint32 flags = 2;\n"
        "        uint32 preempt_count = 3;\n"
        "        int32 pid = 4;\n"
        "    };\n"
        "    CommonFileds common_fields = 50;\n\n"
        "    oneof event {\n"
    )


def to_camel_case(name):
    return "".join([p.capitalize() for p in name.split("_")])


def fix_field_name(name):
    replace_map = {
        "errno": "error_code",
        "sa_handler": "sig_handler",
        "sa_flags": "sig_flags",
        "new": "seq_new",
    }
    if name in replace_map:
        name = replace_map.get(name)
    return str.lower(name)


class FtraceEventProtoGenerator(FtraceEventCodeGenerator):
    def __init__(self, events_dir, allow_list):
        super().__init__(events_dir, allow_list)

    def generate_code(self):
        proto_file_list = []
        for category in self.grouped_event_formats:
            proto_name = "{}.proto".format(category)
            proto_path = os.path.join(self.output_dir, proto_name)
            proto_file_list.append(proto_path)
            self.generate_event_proto(category, proto_path)

        bundle_proto_head = [Common.proto_common_header]
        bundle_proto_body = [Common.proto_message_head]
        event_category_count = 0
        for category in self.grouped_event_formats:
            event_category_count += 1
            proto_name = "{}.proto".format(category)
            bundle_proto_head.append('import "{}";\n'.format(proto_name))
            for i in range(len(self.grouped_event_formats[category])):
                event_format = self.grouped_event_formats[category][i]
                message_name = "{}_format".format(event_format.name)
                message_type = to_camel_case(message_name)
                message_id = event_category_count * 100 + i
                message_name = fix_field_name(message_name)
                bundle_proto_body.append(
                    "        {} {} = {};\n".format(
                        message_type, message_name, message_id
                    )
                )
        bundle_proto_head.append("\n")
        bundle_proto_body.append("    }\n")
        bundle_proto_body.append("}\n")

        bundle_proto_path = os.path.join(self.output_dir, Common.ftrace_bundle_proto)
        proto_file_list.append(bundle_proto_path)
        Common.logger.info("Generate {} ...".format(bundle_proto_path))
        with open(bundle_proto_path, "w+") as f:
            f.writelines(bundle_proto_head + bundle_proto_body)

        protos_gni_path = os.path.join(self.output_dir, Common.auto_generated_gni)
        Common.logger.info("Generate {} ...".format(protos_gni_path))
        with open(protos_gni_path, "w+") as f:
            f.write(Common.gn_copyright_header)
            # proto sources
            f.write("auto_generated_ftrace_proto_sources = [\n")
            for proto_path in proto_file_list:
                f.write('  "{}",\n'.format(os.path.basename(proto_path)))
            f.write("]\n")

    def generate_event_proto(self, category, proto_path):
        with open(proto_path, "w+") as f:
            f.write(Common.proto_common_header)
            f.write("// category: {}\n".format(category))
            Common.logger.info("Generate {} ...".format(proto_path))

            if category in self.grouped_event_formats:
                for event_format in self.grouped_event_formats[category]:
                    Common.logger.debug(
                        "{}/{}:".format(event_format.category, event_format.name)
                    )
                    message_name = "{}_format".format(event_format.name)
                    message_type = to_camel_case(message_name)
                    device_format_file = "{}/{}/{}/format".format(
                        "/sys/kernel/debug/tracing/events",
                        event_format.category,
                        event_format.name,
                    )

                    field_count = 1
                    f.write("// {}\n".format(device_format_file))
                    f.write("message {} {{\n".format(message_type))
                    for field in event_format.remain_fields:
                        type_string = field.to_proto_type().to_string()
                        field_name = fix_field_name(field.name)
                        Common.logger.debug("    {}: {}".format(field, type_string))
                        f.write(
                            "    {} {} = {};\n".format(
                                type_string, field_name, field_count
                            )
                        )
                        field_count += 1
                    f.write("}\n")
                    f.write("\n")


def main():
    parser = argparse.ArgumentParser(description="FTrace proto code generator.")
    parser.add_argument(
        "-a",
        dest="allow_list",
        required=True,
        type=str,
        help="event allow list file path",
    )
    parser.add_argument(
        "-e", dest="events_dir", required=True, type=str, help="event formats directory"
    )
    parser.add_argument(
        "-o",
        dest="output_dir",
        required=True,
        type=str,
        help="code file output directory",
    )

    args = parser.parse_args(sys.argv[1:])
    events_dir = args.events_dir
    output_dir = args.output_dir
    allow_list = args.allow_list

    generator = FtraceEventProtoGenerator(events_dir, allow_list)
    generator.generate(output_dir)


if __name__ == "__main__":
    main()
